#!/usr/bin/env python
# -*- coding: utf-8 -*-

import click
from ase.io import read
from ase import Atoms
import numpy as np
import copy
from itertools import groupby


from . import fourthorder_core  # type: ignore
from .fourthorder_common import (
    SYMPREC,
    gen_SPOSCAR,
    calc_dists,
    calc_frange,
)


@click.group()
def fourthorder():
    """Fourth-order force-constants for VASP."""
    pass


def read_POSCAR():
    """
    Return all the relevant information contained in a POSCAR file.
    """

    atoms = read("POSCAR", format="vasp")

    nruter = dict()
    nruter["lattvec"] = 0.1 * atoms.get_cell().T

    chemical_symbols = atoms.get_chemical_symbols()
    unique_symbols = []
    numbers = []

    for symbol, group in groupby(chemical_symbols):
        unique_symbols.append(symbol)
        numbers.append(len(list(group)))

    nruter["elements"] = unique_symbols
    nruter["numbers"] = np.array(numbers, dtype=np.intc)

    positions = atoms.get_scaled_positions()
    nruter["positions"] = positions.T
    nruter["types"] = []
    nruter["types"] = np.repeat(
        range(len(nruter["numbers"])), nruter["numbers"]
    ).tolist()

    return nruter


def write_POSCAR(poscar, filename):
    symbols = np.repeat(poscar["elements"], poscar["numbers"]).tolist()

    atoms = Atoms(
        symbols=symbols,
        scaled_positions=poscar["positions"].T,
        cell=poscar["lattvec"].T * 10,
    )

    atoms.write(filename, format="vasp", direct=True)


def normalize_SPOSCAR(sposcar):
    """
    Rearrange sposcar, as generated by gen_SPOSCAR, so that it is in
    valid VASP order, and return the result.
    """
    nruter = copy.deepcopy(sposcar)
    indices = np.array(range(nruter["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1)
    )
    indices = np.rollaxis(indices, 3, 0).flatten().tolist()
    nruter["positions"] = nruter["positions"][:, indices]
    nruter["types"].sort()
    return nruter


def read_forces(filename):
    atoms = read(filename)
    nruter = atoms.get_forces()
    return nruter


def build_unpermutation(sposcar):
    """
    Return a list of integers mapping the atoms in the normalized
    version of sposcar to their original indices.
    """
    indices = np.array(range(sposcar["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1)
    )
    indices = np.rollaxis(indices, 3, 0).flatten()
    return indices.argsort().tolist()


def _validate_cutoff(na, nb, nc):
    if min(na, nb, nc) < 1:
        raise click.BadParameter("na, nb and nc must be positive integers")


def _parse_cutoff(cutoff):
    if cutoff.startswith("-"):
        try:
            nneigh = -int(cutoff)
        except ValueError:
            raise click.BadParameter("invalid cutoff")
        if nneigh == 0:
            raise click.BadParameter("invalid cutoff")
        return nneigh, None
    else:
        try:
            frange = float(cutoff)
        except ValueError:
            raise click.BadParameter("invalid cutoff")
        if frange == 0.0:
            raise click.BadParameter("invalid cutoff")
        return None, frange


def _prepare_calculation(na, nb, nc, cutoff):
    """
    Validate the input parameters and prepare the calculation.
    """
    _validate_cutoff(na, nb, nc)
    nneigh, frange = _parse_cutoff(cutoff)
    print("Reading POSCAR")
    poscar = read_POSCAR()
    print("Analyzing the symmetries")
    symops = fourthorder_core.SymmetryOperations(
        poscar["lattvec"], poscar["types"], poscar["positions"].T, SYMPREC
    )
    print(f"- Symmetry group {symops.symbol} detected")
    print(f"- {symops.translations.shape[0]} symmetry operations")
    print("Creating the supercell")
    sposcar = gen_SPOSCAR(poscar, na, nb, nc)
    print("Computing all distances in the supercell")
    dmin, nequi, shifts = calc_dists(sposcar)
    if nneigh is not None:
        frange = calc_frange(poscar, sposcar, nneigh, dmin)
        print(f"- Automatic cutoff: {frange} nm")
    else:
        print(f"- User-defined cutoff: {frange} nm")
    print("Looking for an irreducible set of fourth-order IFCs")

    return poscar, sposcar, symops, dmin, nequi, shifts, frange, nneigh
